
import re
from llm import  llm
from llm import super_eval
import pandas as pd
from rank_bm25 import BM25Okapi
import os
import json
import hashlib
from functools import wraps





def get_query(question,sample):
    prompt = '''
你需要为用户的问题生成搜索查询语句，这些语句将用来检索数据库中的特定字段。
请确保你的查询语句符合提供的JSON格式，并包含中英文字段名称，结果在```json ```内。
json格式为：
```json
[
  {
    "type": "condition",
    "column_name_en": "EnglishName",
    "column_name_zh": "中文名称"
  },
  {
    "type": "results",
    "column_name_en": "EnglishName",
    "column_name_zh": "中文名称"
  },

]
```
其中 type 分为 `condition` 和 `results`
根据问题的具体内容，仔细提取所有关键查询元素。
如果问题内容与示例相似，请模仿示例的风格进行构造。

示例问题：
''' + '\n'.join(sample)

    messages = [{"role": "system", "content": prompt},
                {"role": "user", "content": question}]

    query_list = super_eval(llm(messages))

    return query_list



def intersection_length(str1, str2):
    """计算两个字符串的交集长度"""
    return len(set(str1) & set(str2))


def lcs_length(str1, str2):
    """计算两个字符串的最长公共子序列（LCS）的长度"""
    m, n = len(str1), len(str2)
    dp = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if str1[i - 1] == str2[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

    return dp[m][n]


def calculate_score(str_query, str_item):
    """计算排序得分"""
    intersection_len = intersection_length(str_query, str_item) # 交集长度
    lcs_len = lcs_length(str_query, str_item) # 交集长度 最长公共子序列长度
    return (intersection_len + lcs_len) / len(str_item)


import jieba.analyse
def split_text(text):
    tokens = re.findall('[\u4e00-\u9fa5A-Za-z]', text)
    tokens.extend(jieba.analyse.textrank(text)[:3])
    return tokens
def sort_list_by_score(str_query, str_list):
    """对列表进行排序并返回排序后的索引和得分"""
    # scores = [(i, calculate_score(str_query, str_item)) for i, str_item in enumerate(str_list)]
    # scores = [(i, fuzz.partial_ratio(str_item,str_query)) for i, str_item in enumerate(str_list)]
    # start = time.time()
    bm25 = BM25Okapi([split_text(i) for i in str_list])
    scores = bm25.get_scores(split_text(str_query))
    scores = [(i,sc) for i,sc in enumerate(scores)]
    # print(time.time() - start)
    sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)
    sorted_ids = [i for i, score in sorted_scores]
    sorted_scores_only = [score for i, score in sorted_scores]
    return sorted_ids, sorted_scores_only



data_tb = pd.read_excel('./data/数据字典.xlsx')
tb_name2key  =  dict(zip(data_tb['表英文'],data_tb['key']))

data = pd.read_excel('./data/数据字典.xlsx', sheet_name='表字段信息')

data = data[['table_name','column_name','column_description']] #,'注释','table_name'
data.dropna(inplace=True)
data.drop_duplicates(subset=['table_name','column_name'], inplace=True)
data = data.to_dict('records')

# all_column_name =  [i['column_name'] for i in data]
# for i in data:
#     # if all_column_name.count(i['column_name']) > 3:
#     i['column_description'] = i['column_description'] + f"({tb_name2key[i['table_name']]})"

tb_cn2desc = {}
for i in data:
    tb_cn2desc[i['table_name']+i['column_name']]=i['column_description']


# 示例
str_list = []
for i in data:
    doc = str(i['column_description'])

    str_list.append(doc)

with open('./data/字段关键词.json',encoding='utf-8') as f:
    for line in f:
        key_words = json.loads(line)
        for key in key_words['key_words']:
            str_list.append('筛选字段'+key)
            data.append({
                'table_name':key_words['table'],
                'column_name':key_words['column'],
                'column_description':tb_cn2desc[key_words['table']+key_words['column']],
                '注释':f'此字段可以按照`{key}`筛选'
            })


from get_sample import get_column_sample
searcher_column = get_column_sample()

def search_columns(query):
    if query[2] == '-':
        query = query[2:]
    top_n = 50
    # sample_n = 5
    li = []
    exis_name = []
    sorted_ids, sorted_scores_only = sort_list_by_score(query, str_list)
    for sid,score in zip(sorted_ids, sorted_scores_only):
        if score > 3:
            if data[sid]['column_description'] not in exis_name:
                exis_name.append(data[sid]['column_description'])
                li.append(data[sid])
            if len(li)>top_n:
                break
    li2 = []

    for i in li:
        li2dt = {"table_name":i["table_name"],"column_name_en":i["column_name"],"column_name_zh":i["column_description"]}
        if "注释" in i:
            li2dt["注释"] = i["注释"]
        li2.append(li2dt)
    del li
    li =li2
    sample = searcher_column.get_sample(query)

    sample = [f"user:{i['question']}\nassistant:```json\n{json.dumps(i['columns'],ensure_ascii=False,indent=4)}\n```" for i in sample]
    prompt = """
### 数据库字段搜索查询生成助手

#### 任务说明

你需要根据用户的问题生成搜索查询语句，这些语句将用于检索数据库中的特定字段。请确保你的查询语句符合提供的JSON格式，并包含中英文字段名称。

#### JSON格式

生成的查询语句应符合以下JSON格式：

```json
[
  {
    "column_name_en": "EnglishName",
    "column_name_zh": "中文名称"
  },
  {
    "column_name_en": "EnglishName",
    "column_name_zh": "中文名称"
  }
]
```


#### 工作流程

1. **接收用户问题**：
   - 用户将提供一个问题，你需要根据问题内容提取所有关键查询元素。

2. **字段匹配**：
   - 根据问题内容和提供的关键词，匹配适合的字段。
   - 如果问题内容与示例相似，请模仿示例的风格进行构造。
   - 区分细节，如 `概念` 和 `行业`

3. **生成查询JSON**：
   - 根据匹配的字段，生成符合JSON格式的查询语句。

#### 提示

- 必须包含所有的查询元素。
- 如果匹配的字段和示例中有合适的字段，请使用其中的字段名称。
- 如果匹配的字段和示例中没有合适的字段，请仿照其风格生成名称。

通过关键词匹配的一些字段为："""+json.dumps(li,ensure_ascii=False,indent=2)+"""
示例问题：
""" + '\n'.join(sample)

    messages = [{"role": "system", "content": prompt},
                {"role": "user", "content": f'请输出`{query}`所有需要使用的字段json：'}]

    query_list = super_eval(llm(messages))

    return query_list


if __name__ == '__main__':
    print(search_columns('A股-600872的全称、A股简称、法人、法律顾问、会计师事务所及董秘是？'))